import pandas as pd
import os
import xml.etree.ElementTree as ET

# 设定你的ACE2005文件地址（Chinese）
ace_path = 'D:/work/NLP/ChineseACE2005/Chinese/'
class arg_ID:
    start=0
    end=0
    ID=0
'''class arg_ID_plus:
    start_1=0
    end_1=0
    start_2 = 0
    end_2 = 0
    ID=0'''
if __name__ == '__main__':
    with open('D:/work/NLP/apf.xml-file-in-ACE2005-to-ann/data_list_test.csv', mode='r', encoding='utf-8') as csv_file: #只用于wl
        rows = csv_file.readlines()
        for row in rows[1:]:
            items = row.replace('\n', '').split(',')
            name = items[0]
            path = os.path.join(ace_path, name)
            sgm_path = (path + '.sgm')

            # print(xml_path)
            # 解析xml文件
            xml_path = (path + '.apf.xml')
            apf_tree = ET.parse(xml_path)
            apf_root = apf_tree.getroot()

            # 获取文件名
            document = apf_root.find('document')
            document_name = document.attrib

            # 创建ann文件
            output = open(path+'.ann', 'w', encoding='utf-8')

            # T的下标，和E的下标。
            # 我们暂时设定实体和论元都是T，触发词是E，便于区分。
            T_index = 1
            E_index = 1
            R_index = 1
            # 调整位置
            entity1 = [1000, ""]
            entity2 = [1000, ""]
            for child_document in document:
                # 找到位置次靠前的标记，因为最靠前的往往是POSTER
                entity = child_document
                for child_entity in entity:
                    if child_entity.tag == 'entity_mention':
                        entity_mention = child_entity
                        for head in entity_mention.iter('head'):
                            for entity_charseq in head:  # 遍历 charseq 子节点
                                if entity1[0] > int(entity_charseq.attrib['START']):
                                    entity1[0] = int(entity_charseq.attrib['START'])
                                    entity1[1] = entity_charseq.text
                            for entity_charseq in head:  # 遍历 charseq 子节点
                                if entity2[0] > int(entity_charseq.attrib['START']) and int(entity_charseq.attrib['START']) != entity1[0]:
                                    entity2[0] = int(entity_charseq.attrib['START'])
                                    entity2[1] = entity_charseq.text

            txt = open(sgm_path, 'r', encoding='utf-8')
            # 找到txt中最靠前的标记，计算需调整距离
            pos = 0
            ifHeadline = 0
            turndis = 0
            for line in txt:
                if "<HEADLINE>" in line:
                    ifHeadline = 1
                if "</HEADLINE>" in line:
                    ifHeadline = 0
                if entity2[1] in line:
                    if not ifHeadline:
                        pos = pos + line.index(entity2[1])
                        break
                for ch in line:
                    pos += 1
            dis = int(pos - entity2[0]) - turndis  # 要调整的距离

            #调整删除空格和换行后的位置
            txt = open(sgm_path, 'r', encoding='utf-8')
            txtpos = []  # txtpos[i]存放i位置之前要调整的距离
            add = 0
            flag = 0
            beginflag = 0  # 从<POSTDATE>开始计算
            txtpos.append(0)
            for line in txt:
                print(add)
                if line=="\n":
                    txtpos.append(add)
                    continue
                for ch in line:
                    if beginflag == 0:
                        txtpos.append(add)
                        continue
                    if flag:
                        flag = 0
                        add += 1
                        if ch!='”' and ch!='"':
                            txtpos[-1] += 1
                    if ch == '\n' or ch == ' ':
                        add -= 1
                    elif ch == '。':
                        flag = 1
                    txtpos.append(add)
                if "</POSTDATE>" in line:
                    beginflag = 1


            arg_list=[]
            arg_list_num=0
            ''' arg_list_plus = []
            arg_list_num_plus = 0'''
            for child_document in document:
                # 提取实体

                '''if child_document.tag == 'entity':
                    entity = child_document
                    entity_type = entity.attrib.get('TYPE')
                    for child_entity in entity:
                        if child_entity.tag == 'entity_mention':
                            entity_mention = child_entity
                            entity_mention_ID = 'T' + str(T_index)
                            T_index += 1
                            for head in entity_mention.iter('head'):
                                for entity_charseq in head:  # 遍历 charseq 子节点
                                    entity_mention_start = str(dis + int(entity_charseq.attrib['START']))
                                    entity_mention_end = str(dis + int(entity_charseq.attrib['END']) + 1)
                                    # 设置跨行格式
                                    pos = 0
                                    if '\n' in entity_charseq.text:
                                        pos = entity_charseq.text.index('\n')
                                        entity_charseq.text = entity_charseq.text.replace('\n', ' ')
                                        entity_mention_end = str(
                                            int(entity_mention_start) + pos + 1) + ' ' + entity_mention_end
                                        entity_mention_start = entity_mention_start + ' ' + str(
                                            int(entity_mention_start) + pos)
                                        output.write(
                                            entity_mention_ID + '\t' + entity_type + ' ' + entity_mention_start + ';' +
                                            entity_mention_end + '\t' + entity_charseq.text+'\n')
                                    else:
                                        output.write(
                                            entity_mention_ID + '\t' + entity_type + ' ' + entity_mention_start + ' ' +
                                            entity_mention_end + '\t' + entity_charseq.text+'\n')'''

                # 提取事件
                if child_document.tag == 'event':
                    event = child_document
                    event_type = event.attrib.get('SUBTYPE')
                    event_ID = event.attrib.get('ID')
                    event_argument_role = []
                    event_argument_ID = []
                    event_argument = {}
                    for child_event in event:
                        # 提取 event_argument 属性
                        if child_event.tag == 'event_argument':
                            event_argument_role.append(child_event.get('ROLE'))
                            event_argument_ID.append(child_event.get('REFID'))

                        # 提取 trigger
                        elif child_event.tag == 'event_mention':
                            event_mention = child_event
                            event_mention_id = event_mention.attrib.get('ID')
                            for child_event_mention in event_mention:
                                # 提取触发词
                                if child_event_mention.tag == 'anchor':
                                    trigger = child_event_mention  # trigger 的 id 是event_mention 的 id
                                    trigger_ID = 'E' + str(E_index)
                                    E_index += 1
                                    argument_ID = 'T' + str(T_index)
                                    T_index += 1
                                    for trigger_charseq in trigger:
                                        trigger_start = str(dis + int(trigger_charseq.attrib['START']))

                                        newtrigger_start = str(
                                            dis + int(trigger_charseq.attrib['START']) + txtpos[int(trigger_start)])

                                        # 修改trigger_charseq.text中的空格与换行
                                        newtrigger_text = str(trigger_charseq.text).replace(" ", "")
                                        newtrigger_text = newtrigger_text.replace("\n", "")
                                        newtrigger_end = str(int(newtrigger_start)+len(newtrigger_text))

                                        output.write(
                                            argument_ID + '\t' + event_type + ' ' + newtrigger_start + ' ' +
                                            newtrigger_end + '\t' + newtrigger_text + '\n')
                                        output.write(
                                            trigger_ID + '\t' + event_type + ':' + argument_ID + '\n')
                                        '''argument_ID = 'T' + str(T_index)
                                        T_index += 1
                                        output.write(
                                            argument_ID + '\t' + 'trigger' + ' ' + trigger_start + ' ' +
                                            trigger_end + '\t' + trigger_charseq.text + '\n')
                                        trigger_ID = argument_ID'''

                                        # print(
                                        #   trigger_ID + '\t' + event_type + ' ' + trigger_start + ' ' +
                                        #   trigger_end + '\t' + trigger_charseq.text)

                                # 提取 event_mention 中的 argument
                                elif child_event_mention.tag == 'event_mention_argument':
                                    argument = child_event_mention
                                    argument_ID = 'T' + str(T_index)
                                    T_index += 1
                                    argument_role = argument.attrib['ROLE']
                                    argument_type = ""

                                    for argument_charseq in argument.iter('charseq'):
                                        argument_start = str(dis + int(argument_charseq.attrib['START']))
                                        argument_end = str(dis + int(argument_charseq.attrib['END']))
                                        newargument_start = str(
                                            dis + int(argument_charseq.attrib['START']) + txtpos[int(argument_start)])

                                        p=arg_ID()
                                        p.start=argument_start
                                        p.end=argument_end
                                        p.ID=argument_ID

                                        pos = 0
                                        relation_ID = 'R' + str(R_index)
                                        R_index += 1

                                        poss = 0
                                        for i in range(0, arg_list_num):
                                            if p.start == arg_list[i].start and p.end == arg_list[i].end:
                                                argument_ID = arg_list[i].ID
                                                poss = 1
                                                T_index -= 1
                                        if poss == 0:
                                            arg_list.append(p)
                                            arg_list_num += 1
                                            # 找到argument对应的entity_type或value_type或timex2
                                            argument_REFID = argument.attrib['REFID']
                                            typeFlag = 0
                                            for child_document in document:
                                                if child_document.tag == 'entity':
                                                    entity = child_document
                                                    entity_type = entity.attrib.get('TYPE')
                                                    for child_entity in entity:
                                                        if child_entity.tag == 'entity_mention':
                                                            if child_entity.attrib['ID'] == argument_REFID:
                                                                typeFlag = 1
                                                                argument_type = entity_type
                                            for child_document in document:
                                                if child_document.tag == 'value':
                                                    value = child_document
                                                    value_type = value.attrib.get('TYPE')
                                                    for child_value in value:
                                                        if child_value.tag == 'value_mention':
                                                            if child_value.attrib['ID'] == argument_REFID:
                                                                argument_type = value_type
                                            for child_document in document:
                                                if child_document.tag == 'timex2':
                                                    for child_timex2 in child_document:
                                                        if child_timex2.tag == 'timex2_mention':
                                                            if child_timex2.attrib['ID'] == argument_REFID:
                                                                argument_type = "Timex2"

                                            # 修改argument_charseq.text中的空格与换行
                                            newargument_text=str(argument_charseq.text).replace(" ","")
                                            newargument_text=newargument_text.replace("\n","")

                                            newargument_end = str(int(newargument_start) + len(newargument_text))

                                            output.write(
                                                argument_ID + '\t' + argument_type + ' ' + newargument_start + ' ' +
                                                newargument_end + '\t' +newargument_text + '\n')
                                        output.write(
                                            relation_ID + '\t' + argument_role + ' Arg1:' + trigger_ID + ' Arg2:' + argument_ID + '\n')





                ''' #提取关系
                elif child_document.tag == 'relation':
                    relation = child_document
                    relation_type = relation.attrib.get('TYPE')
                    relation_ID = relation.attrib.get('ID')
                    relation_argument_role = []
                    relation_argument_ID = []
                    relation_argument = {}

                    for child_relation in relation:
                        # 提取 relation_argument 属性
                        biubiu=["0","0"]
                        if child_relation.tag == 'relation_argument':
                            relation_argument_role.append(child_relation.get('ROLE'))
                            relation_argument_ID.append(child_relation.get('REFID'))

                        # 提取 relation_mention
                        elif child_relation.tag == 'relation_mention':
                            relation_mention = child_relation
                            relation_mention_id = relation_mention.attrib.get('ID')
                            for child_relation_mention in relation_mention:


                                # 提取 relation_mention 中的 argument
                                if child_relation_mention.tag == 'relation_mention_argument':
                                    argument = child_relation_mention
                                    argument_ID = 'T' + str(T_index)
                                    T_index += 1
                                    argument_role = argument.attrib['ROLE']
                                    if(argument_role=="Arg-1" ):
                                        biubiu[0]=argument_ID
                                    elif(argument_role=="Arg-2"):
                                        biubiu[1]=argument_ID


                                    for argument_charseq in argument.iter('charseq'):
                                        argument_start = str(dis + int(argument_charseq.attrib['START']))
                                        argument_end = str(dis + int(argument_charseq.attrib['END']) + 1)
                                        # 设置跨行格式
                                        pos = 0
                                        if '\n' in argument_charseq.text:
                                            pos = argument_charseq.text.index('\n')
                                            argument_charseq.text = argument_charseq.text.replace('\n', ' ')
                                            argument_end = str(
                                                int(argument_start) + pos + 1) + ' ' + argument_end
                                            argument_start = argument_start + ' ' + str(
                                                int(argument_start) + pos)
                                            output.write(
                                                argument_ID + '\t' + argument_role + ' ' + argument_start + ';' +
                                                argument_end + '\t' + argument_charseq.text + '\n')
                                        else:
                                            output.write(
                                                argument_ID + '\t' + argument_role + ' ' + argument_start + ' ' +
                                                argument_end + '\t' + argument_charseq.text + '\n')

                            relation_ID = 'R' + str(R_index)
                            R_index += 1
                            if(biubiu[0]!="0"):
                                output.write(relation_ID + '\t' + relation_type + ' Arg1:' +biubiu[0]+' Arg2:' + biubiu[1]+'\n')
                           #     output.write(str(len(biubiu))+'\n')
                            else:
                                output.write(relation_ID + '\t' + relation_type + ' Arg1:' + ' Arg2:' + '\n')'''




            output.close()
